rm(list = ls())
library(reshape2)
library(ggplot2)
library(magrittr)
library(stringr)
devtools::load_all()

# 参数设置
k0 <- seq(0.4,16,0.4)
vlast2 <- vlast1 <- rep(20,length(k0))
beta <- .98
delta <-  .1
theta <- .36
numits <- 250
A1 <- 1.75
A2 <- 0.75

# 图形初始化
picdata <- data.frame(k0 = k0, v = vlast1)

# 值函数迭代numits次。在每个k=(0.01，6.2)上以0作为初值开始迭代
v2 <- k2_opt <- v1 <- k1_opt <- numeric(length(k0))
for (i in 1:numits){
  # 寻找每个k点的最优值函数
  for (j in 1:length(k0)){
    # 优化第一个随机状态
    ans <- optimize(valfun_sto, interval = c(0.41,15.99),At=A1,markov=T, pmkv = c(0.9,0.1),
                    kt = k0[j], beta = beta, theta = theta,
                    delta = delta, k0 = k0, p1=p1,p2=p2,vlast1 = vlast1,vlast2 = vlast2, maximum = T)
    v1[j] <- ans$objective
    k1_opt[j] <- ans$maximum
    # 优化第二个随机状态
    ans <- optimize(valfun_sto, interval = c(0.41,15.99),At=A2,markov=T, pmkv = c(0.4,0.6),
                    kt = k0[j], beta = beta, theta = theta,
                    delta = delta, k0 = k0, p1=p1,p2=p2,vlast1 = vlast1,vlast2 = vlast2, maximum = T)
    v2[j] <- ans$objective
    k2_opt[j] <- ans$maximum
  }

  # 每50次迭代存储一下
  if (i %% 50 == 0){
    print(i)
    picdata[,paste('v1_',as.character(i), sep = '')] <- v1
    picdata[,paste('k1_opt',as.character(i), sep = '')] <- k1_opt
    picdata[,paste('v2_',as.character(i), sep = '')] <- v2
    picdata[,paste('k2_opt',as.character(i), sep = '')] <- k2_opt
  }

  # 替换上一次的值函数
  vlast1 <- v1
  vlast2 <- v2
}

# value function
ans <- melt(picdata, id.vars = 'k0')
ggplot(ans[str_detect(ans$variable,'^v1'),], aes(x = k0, y = value, color = variable)) +
  geom_line() + geom_line(data = ans[str_detect(ans$variable,'^v2'),],
                          aes(x = k0, y = value, color = variable), linetype = 2) +
  scale_color_manual(values = c(v1_50 = 'red',v1_100 = 'blue',v1_150 = 'green',
                                v1_200 = 'brown',v1_250 = 'black',
                                v2_50 = 'red',v2_100 = 'blue',v2_150 = 'green',
                                v2_200 = 'brown',v2_250 = 'black')) + theme_bw()
# ggsave('../val_ite_mkv.png')

# policy function
ggplot(ans[str_detect(ans$variable,'k[1-2]_opt250'),], aes(x = k0, y = value, linetype = variable)) +
  geom_line() + labs(x = 'kt', y = 'ktp1') + theme_bw()
# ggsave('../policy_mkv.png')

# 利用政策函数模拟经济中资本存量的变化
simu_k <- data.frame(t = 1:500, k = NA)
simu_k$k[1] <- 1
status <- 1 # 设定马尔可夫过程的初始状态
for (i in 2:nrow(simu_k)) {
  ans <- runif(1)
  if (status == 1){
    if (ans <= 0.9){
      simu_k$k[i] <- signal::interp1(picdata$k0,picdata$k1_opt250,simu_k$k[i-1],'linear')
      status <- 1
    }else {
      simu_k$k[i] <- signal::interp1(picdata$k0,picdata$k2_opt250,simu_k$k[i-1],'linear')
      status <- 2
    }
  } else if(status == 2){
    if (ans <= 0.4){
      simu_k$k[i] <- signal::interp1(picdata$k0,picdata$k1_opt250,simu_k$k[i-1],'linear')
      status <- 1
    }else {
      simu_k$k[i] <- signal::interp1(picdata$k0,picdata$k2_opt250,simu_k$k[i-1],'linear')
      status <- 2
    }
  }
}
ggplot(data = simu_k, aes(x = t, y = k)) + geom_line() + theme_bw()
# ggsave('../simu_k_mkv.png')
